#include "EnemyDetector.h"
#include"../General/opencv_extended.h"
#include"../General/numeric_rm.h"
#include<chrono>

using namespace std;
using namespace cv;
using namespace cv::ml;


namespace rm
{
#ifdef USE_SVM
Ptr<SVM> svm = StatModel::load<SVM>("/home/nuc/Robomaster2022/Armor/hog_svm_sentry2.xml");
HOGDescriptor hog(Size(48, 48), Size(12, 12), Size(6, 6), Size(6, 6), 9);
#endif


void FourPointsDetector::init()
{
    // load model
    String model = "/home/nuc/Robomaster2022/Armor/vggregnet.onnx";
    net = dnn::readNetFromONNX(model);
}

int FourPointsDetector::detect(cv::Mat image)
{
    auto start = chrono::high_resolution_clock::now();

    // load image data
    float scale = 1.0 / 256.0; //normalize?
    Scalar mean = Scalar(0, 0, 0);
    int h = image.size().height;
    int w = image.size().width;

    Mat blob;
    dnn::blobFromImage(image, blob, scale, Size(W, H), mean, true, false);

    //inference
    net.setInput(blob);
    Mat output = net.forward();

    for(int i = 0; i < 4; i++)
    {
        Point2f p;
        p.x = (float)w * output.at<Vec2f>(0, i)[0] / (float)W;
        p.y = (float)h * output.at<Vec2f>(0, i)[1] / (float)H;
        if(p.x >= w || p.x < 0 || p.y >= h || p.y < 0)
        {
            return NO;
        }
        vertex.push_back(p);  //lightbar vertex  (tl, tr, br, bl)
    }

    //仿照老版代码的评分体系
    Point2f leftLightCenter = (vertex.at(0) + vertex.at(3)) / 2.0;
    Point2f rightLightCenter = (vertex.at(1) + vertex.at(2)) / 2.0;
    float dis = cvex::distance(leftLightCenter, rightLightCenter);


    float leftLightLength = cvex::distance(vertex.at(0), vertex.at(3));
    float rightLightLength = cvex::distance(vertex.at(1), vertex.at(2));

    float LenDiff_ratio = abs(leftLightLength - rightLightLength) / max(leftLightLength, rightLightLength);
    float max_len = max(leftLightLength, rightLightLength);
    float meanLen = (leftLightLength + rightLightLength) / 2;

    float armor_ratio = (dis *(1+LenDiff_ratio)) / max_len;
    type = armor_ratio < armor_small_armor_ratio ? SMALL_ARMOR : BIG_ARMOR;


    float sizeScore = contourArea(vertex) / 15000;//exp(normalized_area);

    float yDiff = abs(leftLightCenter.y - rightLightCenter.y);
    float ratio = dis / meanLen;

    // calculate the rotation score
    // rotation 正比于 ratio_off & y_off
    float ratiOff = (type == BIG_ARMOR) ? max(armor_big_armor_ratio - ratio, float(0)) : max(armor_small_armor_ratio - ratio, float(0));
    float yOff = yDiff / meanLen;
    float rotationScore = ratiOff * ratiOff + yOff * yOff;

    score = sizeScore + rotationScore;

    for(int i = 0; i < 4; i++)
    {
        Point2f p;
        if(i == 0 || i == 3)
        {
            p = 1.5*(vertex.at(i) - leftLightCenter) + leftLightCenter;
        }
        if(i == 1 || i == 2)
        {
            p = 1.5*(vertex.at(i) - rightLightCenter) + rightLightCenter;
        }
        //灯条延长1.5倍，来框住大部分图案
        _vertex.push_back(p);  //armor pattern vertex

    }

    Mat grayImg;
    cvtColor(image, grayImg, COLOR_BGR2GRAY);
    //筛掉不完整的装甲板
    Mat binBrightImg;
    threshold(grayImg, binBrightImg, brightness_threshold, 255, cv::THRESH_BINARY);
    Mat element = cv::getStructuringElement(cv::MORPH_ELLIPSE, cv::Size(3, 3));
    dilate(binBrightImg, binBrightImg, element);

    Rect leftlight = Rect(min(vertex.at(0).x, vertex.at(3).x), min(vertex.at(0).y, vertex.at(3).y),
                          (abs(vertex.at(0).x - vertex.at(3).x) + 1), abs(vertex.at(0).y - vertex.at(3).y));
    Mat leftLightImg = binBrightImg(leftlight);
    const Scalar leftMeanVal = cv::mean(leftLightImg);

    Rect rightlight = Rect(min(vertex.at(1).x, vertex.at(2).x), min(vertex.at(1).y, vertex.at(2).y),
                          (abs(vertex.at(1).x - vertex.at(2).x) + 1), abs(vertex.at(1).y - vertex.at(2).y));

    Mat rightLightImg = binBrightImg(rightlight);
    const Scalar rightMeanVal = cv::mean(rightLightImg);

    //cout<<leftMeanVal[0]<<" "<<rightMeanVal[0]<<endl;

    if(!(leftMeanVal[0] > 150 && rightMeanVal[0] > 150))
    {
        return NO;
    }
    getFrontImg(grayImg);

#ifdef GET_ARMOR_PIC
    _allCnt++;
    if(_allCnt%10 == 1)
    {
        Mat regulatedFrontImg = frontImg;
        if(type == BIG_ARMOR)
        {
            regulatedFrontImg = regulatedFrontImg(Rect(21, 0, 50, 50));
        }
        imwrite("/home/nuc/ArmorTrainingSample/" + to_string((_allCnt-1)/10) + ".bmp", regulatedFrontImg);
    }
//    Mat regulatedFrontImg = frontImg;
//    if(type == BIG_ARMOR)
//    {
//        regulatedFrontImg = regulatedFrontImg(Rect(21, 0, 50, 50));
//    }
//    imwrite("/home/nuc/ArmorTrainingSample/" + to_string(_allCnt) + ".bmp", regulatedFrontImg);
#endif // GET_ARMOR_PIC

    armor_num = isArmorPattern();

    center = (vertex.at(0) + vertex.at(1)+ vertex.at(2)+ vertex.at(3)) / 4.0;

//    if(armor_num == -1)
//    {
//        return NO;
//    }
    if((leftMeanVal[0] > 150 && rightMeanVal[0] <= 150) || (leftMeanVal[0] <= 150 && rightMeanVal[0] > 150))
    {
        return HALF;
    }
    auto end = chrono::high_resolution_clock::now();
    cout << "4points period: " << (static_cast<chrono::duration<double, std::milli>>(end - start)).count() << " ms" << endl;
    return WHOLE;
}

void FourPointsDetector::getFrontImg(const cv::Mat& grayImg)
{
    using cvex::distance;
    const Point2f&
            tl = vertex[0],
            tr = vertex[1],
            br = vertex[2],
            bl = vertex[3];

    int width, height;
    if(type == BIG_ARMOR)
    {
        width = 92;
        height = 50;
    }
    else
    {
        width = 50;
        height = 50;
    }

    Point2f src[4]{Vec2f(tl), Vec2f(tr), Vec2f(br), Vec2f(bl)};
    Point2f dst[4]{Point2f(0.0, 0.0), Point2f(width, 0.0), Point2f(width, height), Point2f(0.0, height)};
    const Mat perspMat = getPerspectiveTransform(src, dst);
    cv::warpPerspective(grayImg, frontImg, perspMat, Size(width, height));
}

int FourPointsDetector::isArmorPattern()
{
    int num = 1;
#ifdef USE_SVM
    Mat regulatedImg;
    if (type == BIG_ARMOR)
    {
        regulatedImg = frontImg(Rect(21, 0, 48, 48));
    }
    else
    {
        regulatedImg = frontImg;
    }

    vector<float> descriptors;
    hog.compute(regulatedImg, descriptors, Size(6, 6));
    int DescriptorDim = descriptors.size();
    Mat _row(1, DescriptorDim, CV_32FC1);
    for (int i = 0; i < DescriptorDim; i++)
    {
        _row.at<float>(0, i) = descriptors[i];
    }
    float result = svm->predict(_row);

    std::cout << "number of armor is " << (int)result << endl;
    //    imshow("...", regulatedImg);
    //num++;

    if(result == 1 || result == 7 || result == 8)
        type = BIG_ARMOR;
    if(result == 2 || result == 3 || result == 4 || result == 5 || result == 9)
        type = SMALL_ARMOR;
    num = result;
#endif
    return num;
}

void EnemyDetector::init()
{
    yolox_dectector.init();
    p4.init();
}

void EnemyDetector::loadImg(const Mat& _srcImg)
{
    srcImg = _srcImg.clone();
}

bool EnemyDetector::detect()
{
    target_vertex.clear();
    int armor_flag;
    armors.clear();
//    robots.clear();
    yolox_dectector.dectect(srcImg);
    objects = yolox_dectector.getObjects();

    if(objects.size() == 0)
    {
#ifdef SHOW_RESULT
        yolox_dectector.draw_objects(srcImg);
        imshow("debug", srcImg);
        waitKey(1);
#endif
        return false;
    }

    for(int i = 0; i < objects.size(); i++)
    {
        if(enemy_color == BLUE && (objects.at(i).label == 0 || objects.at(i).label == 1))
        {
            armor_flag = p4.detect(srcImg(objects.at(i).rect));
            if(armor_flag != NO)
            {
                Armor armor;
                armor.index = i;
                armor.color = BLUE;
                armor.number = p4.armor_num;
                armor.type = p4.type;
                armor.score = p4.score;
                armor.flag = armor_flag;
                armor.rect = objects.at(i).rect;
                armor.vertex = p4.vertex;
                armor.center = p4.center;
                armors.push_back(armor);
            }
        }
        if(enemy_color == RED && (objects.at(i).label == 0 || objects.at(i).label == 1))
        {
            armor_flag = p4.detect(srcImg(objects.at(i).rect));
            if(armor_flag != NO)
            {
                Armor armor;
                armor.index = i;
                armor.color = BLUE;
                armor.number = p4.armor_num;
                armor.type = p4.type;
                armor.score = p4.score;
                armor.rect = objects.at(i).rect;
                armor.vertex = p4.vertex;
                armor.center = p4.center;
                armors.push_back(armor);
            }
        }
        p4.clear();     //p4使用之后不要忘了  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    }

    if(armors.size() == 0)
    {
#ifdef SHOW_RESULT
        yolox_dectector.draw_objects(srcImg);
        imshow("debug", srcImg);
        waitKey(1);
#endif
        return false;
    }

    int target_index = -1;
    float max_score = 0;
    for(int i = 0; i < armors.size(); i++)
    {
        if(armors.at(i).score >= max_score)
        {
            target_index = i;
            max_score = armors.at(i).score;
        }
    }

    if(armors.at(target_index).vertex.size() != 4)
    {
        return false;
    }

    for(int i = 0; i < 4; i++)
    {
        Point2f p;
        p.x = armors.at(target_index).vertex.at(i).x + armors.at(target_index).rect.tl().x;
        p.y = armors.at(target_index).vertex.at(i).y + armors.at(target_index).rect.tl().y;
        target_vertex.push_back(p);
#ifdef SHOW_RESULT
        circle(srcImg, p, 3, Scalar(0, 50 + 50*i, 155), -1);
#endif
    }
    target_center.x = armors.at(target_index).center.x + armors.at(target_index).rect.tl().x;
    target_center.y = armors.at(target_index).center.y + armors.at(target_index).rect.tl().y;

    target_type = armors.at(target_index).type;
    target_num = armors.at(target_index).number;
    target_flag = armors.at(target_index).flag;

#ifdef SHOW_RESULT
    circle(srcImg, target_center, 3, Scalar(255, 255, 0), -1);
    yolox_dectector.draw_objects(srcImg);
    imshow("debug", srcImg);
    waitKey(1);
#endif
    return true;
}

const std::vector<cv::Point2f> EnemyDetector::getTargetVertex() const
{
    return target_vertex;
}

cv::Point2f EnemyDetector::getTargetCenter()
{
    return target_center;
}

int EnemyDetector::getArmorNum()
{
    return target_num;
}

int EnemyDetector::getArmorType()
{
    return target_type;
}

int EnemyDetector::getArmorFlag()
{
    return target_flag;
}

Mat EnemyDetector::getDebugImg()
{
    return srcImg;
}

}
